from torch.utils.data import Dataset
from mmdet.datasets.pipelines.compose import Compose


class CBISDataset(Dataset):
    CLASSES = None

    def __init__(self, cfg, mg_dcms):
        self.test_mode = cfg['test_mode']
        self.proposal_file = None
        self.mg_dcms = mg_dcms
        self.img_prefix = cfg['img_prefix']
        pipeline = cfg['pipeline']
        self.pipeline = Compose(pipeline)

    def __len__(self):
        return len(self.mg_dcms)

    def pre_pipeline(self, results):
        results['img_prefix'] = self.img_prefix
        results['seg_prefix'] = None
        results['proposal_file'] = None
        results['bbox_fields'] = []
        results['mask_fields'] = []

    def __getitem__(self, idx):
        if self.test_mode:
            return self.prepare_test_img(idx)

    def prepare_test_img(self, idx):
        img_file = self.mg_dcms[idx]
        results = dict(img_info=img_file)
        self.pre_pipeline(results)
        result = self.pipeline(results)
        return result


def build_dataset(cfg, mg_dcms):
    dataset = CBISDataset(cfg, mg_dcms)
    return dataset